import yaml
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision import transforms
from models import resnet18
from pathlib import Path
import matplotlib.pyplot as plt
import os
import time
import utils

# 设置中文显示
plt.rcParams["font.family"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False  # 正确显示负号

# 创建保存结果的目录
result_dir = Path(__file__).parent / 'results'
model_dir = result_dir / 'models'
os.makedirs(model_dir, exist_ok=True)  # 确保目录存在

# 设备配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用设备: {device}")

# 加载配置文件
config_path = Path(__file__).parent / 'config' / 'resnet18.yaml'
with open(config_path, 'r') as f:
    config = yaml.safe_load(f)
    resnet18_cfg = config.get('resnet18', {})
    learning_rate = resnet18_cfg.get('learning_rate', 0.002)
    epochs = resnet18_cfg.get('epochs', 10)
    batch_size = resnet18_cfg.get('batch_size', 100)
    momentum = resnet18_cfg.get('momentum', 0.9)
    weight_decay = resnet18_cfg.get('weight_decay', 5e-4)

# 数据集路径
dataset_path = Path(__file__).parent / 'dataset'
os.makedirs(dataset_path, exist_ok=True)

### 数据预处理
transform_train = transforms.Compose([
    transforms.RandomCrop(32, padding=4),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])
transform_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

# 加载数据集
trainset = torchvision.datasets.CIFAR10(
    root=dataset_path, 
    train=True, 
    download=True, 
    transform=transform_train
)
trainloader = torch.utils.data.DataLoader(
    trainset, 
    batch_size=batch_size, 
    shuffle=True, 
    num_workers=0  # windows只能设置成0
)

testset = torchvision.datasets.CIFAR10(
    root=dataset_path, 
    train=False, 
    download=True, 
    transform=transform_test
)
testloader = torch.utils.data.DataLoader(
    testset, 
    batch_size=batch_size, 
    shuffle=False, 
    num_workers=0  # windows只能设置成0
)

classes = utils.CLASSES

# 初始化模型、损失函数和优化器
net = resnet18.Resnet18().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(
    net.parameters(), 
    lr=learning_rate, 
    momentum=momentum, 
    weight_decay=weight_decay
)

# 记录训练过程中的指标
train_losses = []
train_accuracies = []
test_accuracies = []
best_test_acc = 0.0  # 记录最佳测试准确率

# 开始训练
print(f"开始训练，共 {epochs} 个epoch，批次大小: {batch_size}，学习率: {learning_rate}")
start_time = time.time()

for epoch in range(epochs):
    print(f"\nEpoch: {epoch + 1}/{epochs}")
    net.train()  # 切换到训练模式
    sum_loss = 0.0
    correct = 0.0
    total = 0.0
    
    # 训练循环
    for i, data in enumerate(trainloader, 0):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)
        
        # 清零梯度
        optimizer.zero_grad()
        
        # 前向传播、计算损失、反向传播、参数更新
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        # 统计损失和准确率
        sum_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()
        
        # 每10个批次打印一次信息，避免输出过多
        if (i + 1) % 10 == 0:
            print('[epoch:%d, iter:%d] Loss: %.03f | Acc: %.3f%% '
                  % (epoch + 1, (i + 1 + epoch * len(trainloader)), 
                     sum_loss / (i + 1), 100. * correct / total))
    
    # 记录当前epoch的训练指标
    avg_train_loss = sum_loss / (i + 1)
    train_acc = 100. * correct / total
    train_losses.append(avg_train_loss)
    train_accuracies.append(train_acc)
    print(f"训练集平均损失: {avg_train_loss:.3f}, 训练准确率: {train_acc:.3f}%")
    
    ### 测试阶段
    net.eval()  # 切换到评估模式
    with torch.no_grad():  # 禁用梯度计算，节省内存和计算资源
        correct = 0
        total = 0
        for data in testloader:
            inputs, labels = data
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = net(inputs)
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
        
        test_acc = 100. * correct / total
        test_accuracies.append(test_acc)
        print(f'测试集准确率: {test_acc:.3f}%')
        
        # 保存最佳模型
        if test_acc > best_test_acc:
            best_test_acc = test_acc
            model_path = model_dir / f"resnet18_best_epoch{epoch+1}_acc{test_acc:.2f}.pth"
            torch.save({
                'epoch': epoch + 1,
                'model_state_dict': net.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'loss': avg_train_loss,
                'test_accuracy': test_acc
            }, model_path)
            print(f"已保存最佳模型至: {model_path}")

# 训练结束，保存最终模型
final_model_path = model_dir / f"resnet18_final_epoch{epochs}_acc{test_accuracies[-1]:.2f}.pth"
torch.save({
    'epoch': epochs,
    'model_state_dict': net.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': train_losses[-1],
    'test_accuracy': test_accuracies[-1]
}, final_model_path)
print(f"\n训练完成，总耗时: {time.time() - start_time:.2f}秒")
print(f"最终模型已保存至: {final_model_path}")
print(f"最佳测试准确率: {best_test_acc:.3f}%")

### 绘制训练曲线
# 1. 绘制损失曲线
plt.figure(figsize=(12, 5))
plt.subplot(1, 2, 1)
plt.plot(range(1, epochs+1), train_losses, marker='o', color='b', label='训练损失')
plt.title('训练损失曲线')
plt.xlabel('Epoch')
plt.ylabel('损失值')
plt.grid(True, linestyle='--', alpha=0.7)
plt.legend()

# 2. 绘制准确率曲线
plt.subplot(1, 2, 2)
plt.plot(range(1, epochs+1), train_accuracies, marker='s', color='b', label='训练准确率')
plt.plot(range(1, epochs+1), test_accuracies, marker='^', color='r', label='测试准确率')
plt.title('准确率曲线')
plt.xlabel('Epoch')
plt.ylabel('准确率 (%)')
plt.grid(True, linestyle='--', alpha=0.7)
plt.legend()

# 调整布局并保存图像
plt.tight_layout()
plot_path = result_dir / 'training_curves.png'
plt.savefig(plot_path, dpi=300, bbox_inches='tight')
print(f"训练曲线已保存至: {plot_path}")
plt.show()
